"""
工作流管理器
负责加载、管理和动态替换工作流模板中的占位符
"""

import json
import os
import random
import logging
from typing import Dict, Any, Optional, List
from pathlib import Path

logger = logging.getLogger(__name__)


class WorkflowManager:
    """
    工作流管理器
    管理ComfyUI工作流模板，支持动态参数替换
    """
    
    def __init__(self):
        self.template_dir = Path(__file__).parent.parent / "workflow_templates"
        self._templates_cache = {}
        self._load_templates()
    
    def _load_templates(self):
        """加载所有工作流模板"""
        try:
            if not self.template_dir.exists():
                logger.warning(f"工作流模板目录不存在: {self.template_dir}")
                return
            
            for template_file in self.template_dir.glob("*.json"):
                template_name = template_file.stem
                try:
                    with open(template_file, 'r', encoding='utf-8') as f:
                        template_data = json.load(f)
                    self._templates_cache[template_name] = template_data
        
                except Exception as e:
                    logger.error(f"加载工作流模板失败 {template_file}: {e}")
        
        except Exception as e:
            logger.error(f"加载工作流模板目录失败: {e}")
    
    def get_available_styles(self) -> List[str]:
        """获取所有可用的工作流风格"""
        return list(self._templates_cache.keys())
    
    def get_template(self, style_name: str) -> Optional[Dict[str, Any]]:
        """获取指定风格的工作流模板"""
        return self._templates_cache.get(style_name)
    
    def create_workflow(
        self, 
        style_name: str, 
        character_trigger: Optional[str], 
        user_prompt: str,
        **kwargs
    ) -> Optional[Dict[str, Any]]:
        """
        创建动态工作流
        
        Args:
            style_name: 工作流风格名称
            character_trigger: 角色触发词
            user_prompt: 用户输入提示词
            **kwargs: 其他可选参数（如seed、steps等）
            
        Returns:
            Dict: 完整的工作流配置，失败时返回None
        """
        try:
            # 获取模板
            template = self.get_template(style_name)
            if not template:
                logger.error(f"未找到工作流模板: {style_name}")
                return None
            
            # 深拷贝模板
            workflow = json.loads(json.dumps(template))
            
            # 替换占位符
            workflow = self._replace_placeholders(workflow, {
                "CHARACTER_TRIGGER": character_trigger or "",
                "USER_PROMPT": user_prompt,
                **kwargs
            })
            
            # 生成随机种子
            if "3" in workflow and "inputs" in workflow["3"]:
                workflow["3"]["inputs"]["seed"] = random.randint(1, 2**32 - 1)
            
            

            
            return workflow
            
        except Exception as e:
            logger.error(f"创建工作流失败: {e}", exc_info=True)
            return None
    
    def _replace_placeholders(self, data: Any, replacements: Dict[str, str]) -> Any:
        """
        递归替换数据中的占位符
        
        Args:
            data: 要处理的数据
            replacements: 替换映射
            
        Returns:
            处理后的数据
        """
        if isinstance(data, str):
            # 替换字符串中的占位符
            for placeholder, value in replacements.items():
                data = data.replace(f"{{{{{placeholder}}}}}", str(value))
            return data
        
        elif isinstance(data, dict):
            # 递归处理字典
            return {k: self._replace_placeholders(v, replacements) for k, v in data.items()}
        
        elif isinstance(data, list):
            # 递归处理列表
            return [self._replace_placeholders(item, replacements) for item in data]
        
        else:
            # 其他类型直接返回
            return data
    
    def get_workflow_info(self, style_name: str) -> Optional[Dict[str, Any]]:
        """
        获取工作流信息
        
        Args:
            style_name: 工作流风格名称
            
        Returns:
            工作流信息字典
        """
        template = self.get_template(style_name)
        if not template:
            return None
        
        # 分析工作流信息
        info = {
            "name": style_name,
            "nodes_count": len(template),
            "has_lora": any(
                node.get("class_type") == "LoraLoader" 
                for node in template.values() 
                if isinstance(node, dict)
            ),
            "image_size": "720x1280",  # 默认尺寸
            "description": self._get_style_description(style_name)
        }
        
        # 提取图像尺寸
        for node in template.values():
            if isinstance(node, dict) and node.get("class_type") == "EmptyLatentImage":
                inputs = node.get("inputs", {})
                width = inputs.get("width", 720)
                height = inputs.get("height", 1280)
                info["image_size"] = f"{width}x{height}"
                break
        
        # 提取LoRA信息
        lora_list = []
        for node in template.values():
            if isinstance(node, dict) and node.get("class_type") == "LoraLoader":
                inputs = node.get("inputs", {})
                lora_name = inputs.get("lora_name", "")
                strength = inputs.get("strength_model", 1.0)
                if lora_name:
                    lora_list.append({
                        "name": lora_name,
                        "strength": strength
                    })
        
        info["lora_models"] = lora_list
        return info
    
    def _get_style_description(self, style_name: str) -> str:
        """获取风格描述"""
        descriptions = {
            "AS1-G001": "角色专用版本，包含特定角色LoRA和米山舞风格，适合角色定制",
            "AS1": "综合平涂画风，融合多种LoRA增强效果，适合生成精细的动漫角色",
            "AS2": "米山舞画师风格，结合Niji和星光效果，更注重艺术感和视觉冲击",
            "AS3": "简化基础版本，快速生成标准动漫风格"
        }
        return descriptions.get(style_name, "动漫风格工作流")
    
    def validate_workflow(self, workflow: Dict[str, Any]) -> bool:
        """
        验证工作流的有效性
        
        Args:
            workflow: 工作流配置
            
        Returns:
            bool: 是否有效
        """
        try:
            # 检查必要的节点
            required_nodes = ["3", "4", "5", "6", "7", "8"]  # KSampler, Checkpoint, Latent, CLIP等
            
            for node_id in required_nodes:
                if node_id not in workflow:
                    logger.warning(f"工作流缺少必要节点: {node_id}")
                    return False
            
            # 检查节点结构
            for node_id, node_data in workflow.items():
                if not isinstance(node_data, dict):
                    logger.warning(f"节点数据格式错误: {node_id}")
                    return False
                
                if "class_type" not in node_data:
                    logger.warning(f"节点缺少class_type: {node_id}")
                    return False
            
            # 检查是否还有未替换的占位符
            workflow_str = json.dumps(workflow)
            if "{{" in workflow_str and "}}" in workflow_str:
                logger.warning("工作流中还有未替换的占位符")
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"验证工作流时出错: {e}")
            return False
    
    def reload_templates(self):
        """重新加载模板"""

        self._templates_cache.clear()
        self._load_templates()


# 创建全局工作流管理器实例
workflow_manager = WorkflowManager() 